import sys
sys.path.append("..")  # 也可以这样
import os
import time
from utils.cloud_utils import main_upload_file_to_aliyun, main_upload_file_to_qiniu
from external.tools import findCosineDistance_V2, 向量均值
from external.tools import 读取图片, read_image
from config import scale_factor
import tempfile
from external.tools import read_image
from config import kLocUserImagesPath
from utils.many_face_fenxi import crop_and_save_face_image
from utils.image_utils import process_user_image
from PIL import Image
from config import kLocUserImagesPath, kSourceVideosFrame
import requests
import numpy as np
from typing import Any, List, Callable, Optional
from loguru import logger
from celery.signals import worker_init
from celery.signals import worker_process_init
import threading
from uuid import uuid4
import cv2
from config import similarity_threshold
from Coordinator import FaceModels
from . import app

from utils.cloud_utils import notify_cloud_v2



Frame = List[List[Any]]

# THREAD_LOCK = threading.Lock()
# F = None


@worker_init.connect
def init_worker(**kwargs):
    global F
    try:
        F = FaceModels()

        # 加入这些
        import os
        from config import BASE_DIR
        # 这的2张图片随便换下都可以,可能越小初始化越快吧
        data = {
            "target_file": os.path.join(BASE_DIR, "img/1698143160340.jpg"),
            "user_image": os.path.join(BASE_DIR, "img/1698212401438.jpg"),
            "output_file_path": f'1.jpg'
        }

        start_V1(**data)
        logger.info("F 模型初始化完毕")

    except Exception as e:
        import traceback
        traceback_info = traceback.format_exc()
        print(traceback_info)
        print(f"Error initializing F: {e}")  # 或使用日志工具进行记录


@app.task
def start_V1(user_image, target_file, output_file_path, GFPGANAmount=100):
    data = {
        "target_image": target_file,
        "source_images": [
            {
                # 第二张目标人脸的图片
                "index": 0,
                "face_image_urls": [user_image]
            }
        ]
    }
    return start_V2(data, output_file_path, GFPGANAmount)


# 把这个封装成Celery就可以了!
@app.task
def start_V2(data, output_file_path, GFPGANAmount=100):
    '''
    GFPGANAmount: 人脸修复的程度,100就是修复得最多
    # 如果要视频换脸的话 核心就改这个逻辑
    推荐传递的参数,可以多张图片进行融合
    data = {
    "target_image": "img/WX20230831-151640@2x.png",
    "source_images": [
        {
            # 第一张目标人脸的图片
            "index": 0,
            "face_image_urls": ["img/1698212401438.jpg"]
        },
        {
            # 第二张目标人脸的图片
            "index": 1,
            "face_image_urls": ["img/1693484879724.jpg"]
        }
    ]
}
    '''
    target_image = read_image(data["target_image"])  # 读取图片的frame
    img = target_image
    target_faces = F.从图片获取人脸数据(target_image, max_num=8)
    if target_faces:  # 如果从图片中读取到人脸,才走换脸的流程,否则就直接保存
        # 读取源图片
        for source_image in data["source_images"]:
            target_faces_index = source_image["index"]
            source_image_face_e_list = []   # 把需要换脸的list就放在这里
            for source_image_sub_images_url in source_image["face_image_urls"]:
                # 从源图片获取人脸的信息(这个是用户传的图片)
                source_image_face_data = F.从图片获取人脸数据(
                    read_image(source_image_sub_images_url), max_num=1)
                if source_image_face_data:
                    source_image_face_e = source_image_face_data[0].embedding
                    target_face_embedding = target_faces[target_faces_index].embedding
                    # 比较人脸的相似度
                    sim = findCosineDistance_V2(
                        target_face_embedding, source_image_face_e)
                    if sim > 0.9:
                        pass
                    else:
                        source_image_face_e_list.append(source_image_face_e)
            s_e = 向量均值(source_image_face_e_list)
            # 执行了一次换脸的逻辑
            img = F.swap_core(img, target_faces[target_faces_index].kps, s_e,
                              target_faces[target_faces_index].bbox, GFPGANAmount)  # 只要有向量就可以了
    # 保存图片
    cv2.imwrite(output_file_path, img)
    return output_file_path


@app.task(autoretry_for=(Exception,), retry_backoff=1)
def start_V3(video_image_path, face_pairs_list, output_file_path, GFPGANAmount=100, similarity_threshold=0.35):
    """
    video_image_path: 视频的某一帧图片
    face_pairs_list: 人脸对列表,[(目标人脸向量,源人脸向量),...)]
    output_file_path: 输出的图片路径
    GFPGANAmount: 高清修复的重回,100就是修复得最多
    similarity_threshold: 相似度的阈值,只有相似到一定地步以后,才会进行更换
    """
    frame = read_image(video_image_path)
    frame_faces = F.从图片获取人脸数据(frame, max_num=8)
    if frame_faces:
        # 将列表转换为 float 类型的 NumPy 数组
        face_pairs = [(np.array(target_vector, dtype=np.float32), np.array(source_vector, dtype=np.float32))
                      for target_vector, source_vector in face_pairs_list]
        logger.debug(f"frame_faces: {len(face_pairs)}")
        for frame_face in frame_faces:
            frame_face_vector = frame_face['embedding']  # 直接使用预提取的特征向量
            for target_face_vector, source_face_vector in face_pairs:
                logger.debug(f"frame_face: 开始换脸")
                # 和目标人脸相似度比较高的情况下才换脸
                sim = findCosineDistance_V2(
                    frame_face_vector, target_face_vector)
                # 只要有一定的相似度就直接给他换就可以!
                #! 这里有调整过
                print(
                    f"相似度:{sim},阈值:{similarity_threshold},是否相似:{sim > similarity_threshold}")
                if sim > similarity_threshold:
                    # 执行换脸操作
                    frame = F.swap_core(
                        frame, frame_face['kps'], source_face_vector, frame_face['bbox'], GFPGANAmount)
                    logger.debug(f"不相似,开始换脸: {sim}")
                else:
                    logger.debug(f"不换脸: {sim}")

    cv2.imwrite(output_file_path, frame)
    return output_file_path


@app.task(autoretry_for=(Exception,), retry_backoff=1)
def start_V3_1(video_image_path, face_pairs_list, output_file_path, GFPGANAmount=100, similarity_threshold=0.45):
    frame = read_image(video_image_path)
    frame_faces = F.从图片获取人脸数据(frame, max_num=8)
    if frame_faces:
        # 将列表转换为 float 类型的 NumPy 数组
        face_pairs = [(np.array(target_vector, dtype=np.float32), np.array(source_vector, dtype=np.float32))
                      for target_vector, source_vector in face_pairs_list]
        logger.debug(f"frame_faces: {len(face_pairs)}")
        for frame_face in frame_faces:
            frame_face_vector = frame_face['embedding']  # 直接使用预提取的特征向量
            for target_face_vector, source_face_vector in face_pairs:
                # 不做相似度,直接开始就进行换脸
                logger.debug(f"frame_face: 开始换脸")
                frame = F.swap_core(
                    frame, frame_face['kps'], source_face_vector, frame_face['bbox'], GFPGANAmount)
    cv2.imwrite(output_file_path, frame)
    return output_file_path


# 读取任意的在线图片的地址,成为OpenCV图像
def read_image(source_path_or_url: str) -> Frame:
    if source_path_or_url.startswith(("http://", "https://")):
        # 如果路径是一个URL
        response = requests.get(source_path_or_url)
        response.raise_for_status()  # 如果请求出错，引发异常

        # 将图像从原始数据转换为OpenCV图像格式
        image_array = np.asarray(bytearray(response.content), dtype=np.uint8)
        image = cv2.imdecode(image_array, cv2.IMREAD_COLOR)

    else:
        # 如果路径是一个本地文件
        image = cv2.imread(source_path_or_url)

    return image


def convert_floats(data):
    if isinstance(data, list):
        return [convert_floats(item) for item in data]
    elif isinstance(data, dict):
        return {key: convert_floats(value) for key, value in data.items()}
    elif isinstance(data, np.float32):
        return float(data)
    elif isinstance(data, np.int64):
        return int(data)
    else:
        return data


# @app.task(autoretry_for=(Exception,), retry_kwargs={'max_retries': 2, 'countdown': 1})
@app.task
def detect_faces(user_img_path):
    source_frame = read_image(user_img_path)
    res_datas = F.从图片获取人脸数据(source_frame, max_num=8)
    res_datas = sorted(
        res_datas, key=lambda face: face['bbox'][0])  # 先从左往右排序一下

    item_data = []
    # 提取需要的人脸数据
    for i, data in enumerate(res_datas):
        # 这里需要的数据就可以适当的多一些
        item = {
            "index_id": i,
            "bbox": list(data['bbox']),
        }
        item_data.append(item)

    return convert_floats(item_data)


@app.task
def get_best_face_urls(video_url):
    urls = F.从视频中获取最佳人脸数(video_url=video_url)
    new_urls = []
    for url in urls:
        if "aliyuncs.com/" in url:
            url = url.split("aliyuncs.com/")[-1]
            new_urls.append(url)
            # 把人脸封了
    return new_urls

# 人脸对解析,人脸向量解析


@app.task
def preprocess_source_images(data):
    return F.preprocess_source_images(data)

# 人脸对解析,人脸向量解析第二版


@app.task
def preprocess_source_images_v2(data, task_id=None):
    try:
        return F.preprocess_source_images_v2(data)
    except Exception as e:
        import traceback
        traceback.print_exc()
        notify_cloud_v2(result={},
                        task_id=task_id,
                        status=250,
                        error_message="人脸预处理失败"
                        )
        raise e
        # return {}


# ======================= 新版的接口 ========================================================================================================
@app.task(bind=True)
def get_best_face_urls_v2(self, video_url, oss_name="ali", notify=True):
    try:
        from datetime import datetime
        start_time = datetime.now()
        from utils.cloud_utils import notify_cloud_v2
        logger.debug(video_url)
        urls = F.从视频中获取最佳人脸数(video_url=video_url, oss_name=oss_name)
        new_urls = []
        for url in urls:
            logger.debug(url)
            if "aliyuncs.com/" in url:
                url = url.split("aliyuncs.com/")[-1]
            new_urls.append(url)

        end_time = datetime.now()
        logger.info(f"get_best_face_urls_v2耗时: {end_time - start_time}")
        duration = (end_time - start_time).total_seconds()
        result = new_urls  # 这是要发送的任务结果

        status = 100 if result else 250
        error_message = ""

    except Exception as e:
        status = 250
        error_message = str(e)
        result = []

    # TODO: 每一步都要进行通知!
    finally:
        if notify:
            # 这是需要传递的参数
            notify_cloud_v2(result=result,
                            task_id=self.request.id,
                            status=status,
                            error_message=error_message
                            )
    return {"result": new_urls, "duration": duration}


@app.task
def get_images_faces(user_url, oss_name):
    current_directory = os.getcwd()
    with tempfile.TemporaryDirectory(dir=current_directory) as output_folder:
        try:
            loc_user_path_path = process_user_image(
                user_url, user_images_name=kLocUserImagesPath, max_dimension=350)
            target_image = read_image(loc_user_path_path)
            target_faces = F.从图片获取人脸数据(target_image, max_num=8)
            logger.debug(f"获取到人脸{len(target_faces)}个")
            face_index = 0
            res_urls = []
            for i, face in enumerate(target_faces):
                logger.debug(f"图片的裁剪系数: {scale_factor}")
                res_loc_url = crop_and_save_face_image(loc_user_path_path,
                                                       face, output_folder, i, scale_factor=scale_factor)

                # 判断要上传到阿里云还是七牛
                if oss_name == "ali":
                    res_urls.append(main_upload_file_to_aliyun(res_loc_url))
                else:
                    res_urls.append(main_upload_file_to_qiniu(res_loc_url))
        finally:
            # 不管怎么样,要删除掉图片
            from tybase.tools.image_utils import delete_file_or_directory
            delete_file_or_directory(loc_user_path_path)
        return res_urls
